{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"CUDA_DEVICE_ORDER\"]=\"PCI_BUS_ID\"\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"]=\"0\"\n",
    "import tensorflow as tf\n",
    "gpu_options = tf.GPUOptions(allow_growth=True)\n",
    "sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))\n",
    "\n",
    "from PIL import Image\n",
    "import os.path as osp\n",
    "import glob  \n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from tqdm import tqdm\n",
    "import cv2\n",
    "from keras.models import load_model\n",
    "from helper import *\n",
    "\n",
    "mod = load_model('mod.h5')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def edge_detecton(path):\n",
    "    '''\n",
    "    get sketch\n",
    "    '''\n",
    "    from_mat = cv2.imread(path)\n",
    "    width = float(from_mat.shape[1])\n",
    "    height = float(from_mat.shape[0])\n",
    "    new_width = 0\n",
    "    new_height = 0\n",
    "    if (width > height):\n",
    "        from_mat = cv2.resize(from_mat, (512, int(512 / width * height)), interpolation=cv2.INTER_AREA)\n",
    "        new_width = 512\n",
    "        new_height = int(512 / width * height)\n",
    "    else:\n",
    "        from_mat = cv2.resize(from_mat, (int(512 / height * width), 512), interpolation=cv2.INTER_AREA)\n",
    "        new_width = int(512 / height * width)\n",
    "        new_height = 512\n",
    "    from_mat = from_mat.transpose((2, 0, 1))\n",
    "    light_map = np.zeros(from_mat.shape, dtype=np.float)\n",
    "    for channel in range(3):\n",
    "        light_map[channel] = get_light_map_single(from_mat[channel])\n",
    "    light_map = normalize_pic(light_map)\n",
    "    light_map = resize_img_512_3d(light_map)\n",
    "    line_mat = mod.predict(light_map, batch_size=1)\n",
    "    line_mat = line_mat.transpose((3, 1, 2, 0))[0]\n",
    "    line_mat = line_mat[0:int(new_height), 0:int(new_width), :]\n",
    "    \n",
    "    line_mat = np.amax(line_mat, 2)\n",
    "\n",
    "    sketchKeras = show_active_img_and_save_denoise('sketchKeras', line_mat, 'sketchKeras.jpg')\n",
    "\n",
    "    return sketchKeras"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "source_data_path = \"original image path\"#\"/data4/wangpengxiao/danbooru2017/original\"\n",
    "source_img_path = glob.glob(osp.join(source_data_path,'*/*.jpg'))\n",
    "source_img_path += glob.glob(osp.join(source_data_path,'*/*.png'))\n",
    "source_img_path = sorted(source_img_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#get sketch\n",
    "sketch_path = \"sketch save path\"#\"/data4/wangpengxiao/danbooru2017/original_sketch\"\n",
    "os.mkdir(sketch_path)\n",
    "for path in  tqdm(source_img_path):\n",
    "    sketch_img = edge_detecton(path)\n",
    "    cv2.imwrite(osp.join(sketch_path, osp.basename(path)), sketch_img)    "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
